package org.xunche.agent;
import org.objectweb.asm.*;
import org.objectweb.asm.commons.AdviceAdapter;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

public class TimeAgentByJava {
    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new TimeClassFileTransformer());
    }
    private static class TimeClassFileTransformer implements ClassFileTransformer {
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
            if (className.startsWith("java") || className.startsWith("jdk") || className.startsWith("javax") || className.startsWith("sun") || className.startsWith("com/sun")|| className.startsWith("org/xunche/agent")) {
                //return null或者执行异常会执行原来的字节码
                return null;
            }
            System.out.println("loaded class: " + className);
            ClassReader reader = new ClassReader(classfileBuffer);
            ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
            reader.accept(new TimeClassVisitor(writer), ClassReader.EXPAND_FRAMES);
            return writer.toByteArray();
        }
    }

    public static class TimeClassVisitor extends ClassVisitor {
        public TimeClassVisitor(ClassVisitor classVisitor) {
            super(Opcodes.ASM6, classVisitor);
        }
        @Override
        public MethodVisitor visitMethod(int methodAccess, String methodName, String methodDesc, String signature, String[] exceptions) {
            MethodVisitor methodVisitor = cv.visitMethod(methodAccess, methodName, methodDesc, signature, exceptions);
            return new TimeAdviceAdapter(Opcodes.ASM6, methodVisitor, methodAccess, methodName, methodDesc);
        }
    }
    public static class TimeAdviceAdapter extends AdviceAdapter {
        private String methodName;
        protected TimeAdviceAdapter(int api, MethodVisitor methodVisitor, int methodAccess, String methodName, String methodDesc) {
            super(api, methodVisitor, methodAccess, methodName, methodDesc);
            this.methodName = methodName;
        }
        @Override
        protected void onMethodEnter() {
            //在方法入口处植入
            if ("<init>".equals(methodName)|| "<clinit>".equals(methodName)) {
                return;
            }
            String className = getClass().getName();
            String s = className + "." + methodName;
            TimeHolder.start(s);
        }
        @Override
        protected void onMethodExit(int i) {
            //在方法出口植入
            if ("<init>".equals(methodName) || "<clinit>".equals(methodName)) {
                return;
            }
            String className = getClass().getName();
            String s = className + "." + methodName;
            long cost = TimeHolder.cost(s);
            System.out.println(s + ": " + cost);
        }
    }
}